package com.example.edu.Utils;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.FieldValue;
import co.elastic.clients.elasticsearch.core.*;
import co.elastic.clients.elasticsearch.core.bulk.BulkOperation;
import co.elastic.clients.elasticsearch.core.search.Hit;
import com.example.edu.Dto.ESContentDto;
import com.example.edu.Entity.EsDocument;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.annotation.PreDestroy;
import javax.annotation.Resource;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Component
public class EsUtil {
    @Resource
    ElasticsearchClient client;
    private static final Logger logger = LoggerFactory.getLogger(EsUtil.class);

    @PreDestroy
    public void close() {
        try {
            if (client != null)
                client._transport().close();
        } catch (IOException e){
            logger.error(e.getMessage());
        }
    }

    public void insertDocument(String documentId, EsDocument document){
        try {
            IndexRequest<EsDocument> indexRequest = IndexRequest.of(i -> i
                    .index("document")
                    .id(documentId)
                    .document(document)
            );
            client.index(indexRequest);
        } catch (Exception e){
            logger.error(e.getMessage());
            throw new RuntimeException(e);
        }
    }

    public void insertDocuments(List<String> documentIds, List<EsDocument> documents){
        try {
            // 构建 BulkRequest
            BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
            for (int i = 0; i < documentIds.size(); i++) {
                String documentId = documentIds.get(i);
                EsDocument document = documents.get(i);

                // 将 IndexRequest 封装为 BulkOperation，并添加到 BulkRequest 中
                BulkOperation bulkOperation = BulkOperation.of(b -> b.index(ir -> ir
                        .index("document")
                        .id(documentId)
                        .document(document)
                ));

                // 将 BulkOperation 添加到请求中
                bulkRequestBuilder.operations(bulkOperation);
            }
            // 执行 Bulk 请求
            BulkRequest bulkRequest = bulkRequestBuilder.build();
            BulkResponse bulkResponse = client.bulk(bulkRequest);

            // 检查响应中的错误信息
            if (bulkResponse.errors()) {
                bulkResponse.items().forEach(item -> {
                    if (item.error() != null) {
                        logger.error("Error for document " + item.index() + ": " + item.error().reason());
                    }
                });
                throw new RuntimeException("Bulk insert encountered errors.");
            }

        } catch (IOException e) {
            logger.error(e.getMessage());
            throw new RuntimeException("Error during bulk insert",e);
        }
    }

    public void insertDocuments(List<EsDocument> documents){
        try {
            // 构建 BulkRequest
            BulkRequest.Builder bulkRequestBuilder = new BulkRequest.Builder();
            for (EsDocument document : documents) {
                // 将 IndexRequest 封装为 BulkOperation，并添加到 BulkRequest 中
                BulkOperation bulkOperation = BulkOperation.of(b -> b.index(ir -> ir
                        .index("document")
                        .document(document)
                ));

                // 将 BulkOperation 添加到请求中
                bulkRequestBuilder.operations(bulkOperation);
            }
            // 执行 Bulk 请求
            BulkRequest bulkRequest = bulkRequestBuilder.build();
            BulkResponse bulkResponse = client.bulk(bulkRequest);

            // 检查响应中的错误信息
            if (bulkResponse.errors()) {
                bulkResponse.items().forEach(item -> {
                    if (item.error() != null) {
                        logger.error("Error for document " + item.index() + ": " + item.error().reason());
                    }
                });
                throw new RuntimeException("Bulk insert encountered errors.");
            }

        } catch (IOException e) {
            logger.error(e.getMessage());
            throw new RuntimeException("Error during bulk insert", e);
        }
    }

    // 获取文档
    public EsDocument getDocument(String esDocumentId) throws IOException {
        GetRequest getRequest = GetRequest.of(g -> g
                .index("document")
                .id(esDocumentId)
        );

        return client.get(getRequest, EsDocument.class).source();
    }

    // 更新文档
    public void updateDocument(String documentId, EsDocument document) throws IOException {
        UpdateRequest<Object, Object> updateRequest = UpdateRequest.of(u -> u
                .index("document")
                .id(documentId)
                .doc(document)
        );
        client.update(updateRequest, EsDocument.class);
    }

    // 删除文档
    public void deleteDocument(String documentId) throws IOException {
        DeleteRequest deleteRequest = DeleteRequest.of(d -> d
                .index("document")
                .id(documentId)
        );
        client.delete(deleteRequest);
    }

    public void deleteDocuments(List<String> documentIds) {
        try {
            // 将 List<String> 转换为 List<FieldValue>
            List<FieldValue> fieldValues = documentIds.stream()
                    .map(FieldValue::of)
                    .collect(Collectors.toList());  // 兼容 Java 8+
            // 构建 DeleteByQueryRequest
            DeleteByQueryRequest request = new DeleteByQueryRequest.Builder()
                    .index("document")
                    .query(q -> q.terms(t -> t
                            .field("document_id")
                            .terms(v -> v.value(fieldValues))
                    ))
                    .build();
            // 执行删除请求
            DeleteByQueryResponse response = client.deleteByQuery(request);
        } catch (IOException e) {
            logger.error("Error during bulk delete: {}", e.getMessage(), e);
            throw new RuntimeException("Error during bulk delete", e);
        }
    }




    // 查询文档
    public void searchDocuments(String field, String value) throws IOException {
        SearchRequest searchRequest = SearchRequest.of(s -> s
                .index("document")
                .query(q -> q
                        .match(m -> m
                                .field(field)
                                .query(value)
                        )
                )
        );

        SearchResponse<EsDocument> response = client.search(searchRequest, EsDocument.class);

        for (Hit<EsDocument> hit : response.hits().hits()) {
            System.out.println("Document ID: " + hit.id() + ", Source: " + hit.source());
        }
    }

    public List<ESContentDto> searchByEmbedding(List<Double> queryEmbedding,int num) throws IOException {
        List<Float> queryEmbeddingFloat = queryEmbedding.stream().map(Double::floatValue).collect(Collectors.toList());
                SearchRequest searchRequest = SearchRequest.of(s -> s
                .index("document")
                .query(q -> q
                        .knn(knn -> knn
                                .field("embedding") // 向量字段
                                .queryVector(queryEmbeddingFloat) // 查询的向量
                                .k(num) // 返回前num个最相似的文档
                                .numCandidates(10000) // 设定候选文档数量（可调整）
                        )
                )
        );

        // 执行搜索
        SearchResponse<EsDocument> response = client.search(searchRequest, EsDocument.class);

        List<ESContentDto> r = new ArrayList<>();
        // 处理响应结果
        for (Hit<EsDocument> hit : response.hits().hits()) {
            double similarityScore = hit.score() != null ? hit.score() : 0.0; // 获取相似度分数
            String content = hit.source().getContent() != null ? hit.source().getContent() : "";
            r.add(new ESContentDto(content,similarityScore));
        }
        return r;
    }
}
